#include <tinx/debug.h>
#include <tinx/memory.h>
#include <tinx/arena.h>
#include <tinx/string.h>
#include <tinx/stdlib.h>

static arena_descriptor_t descriptors[DESC_COUNT]; // 内核的内存块描述符目录

static block_t *arena_block(int idx, arena_t *arena, u16 size)
{
    u32 addr = (u32)(arena + 1);
    addr += idx * size;
    block_t *block = (block_t *)addr;
    return block;
}

static arena_t *block_arena(block_t *block)
{
    u32 addr = (u32)block;
    // 页的开头就是 arena 结构
    addr &= ~0xfff;
    // 得到 arena
    arena_t *arena = (arena_t *)addr;
    return arena;
}

void *kmalloc(size_t size)
{
    if (size >= 1024)
    {
        // 计算需要的页数量
        size += sizeof(arena_t);
        u32 count = div_round_up(size, PAGE_SIZE);

        // 构建 arena
        arena_t *arena = (arena_t *)alloc_kpage(count);
        memset(arena, 0, sizeof(arena_t));
        arena->page = true;
        arena->count = count;

        return (arena + 1);
    }

    arena_descriptor_t *desc = NULL;
    for (size_t i = 0; i < DESC_COUNT; i++)
    {
        if (descriptors[i].size >= size)
        {
            desc = &descriptors[i];
            break;
        }
    }
    if (!desc)
    {
        panic("No arena descriptor size %d!!!\n", size);
    }

    if (list_empty(&desc->free_list))
    {
        // 分配 arena
        arena_t *arena = (arena_t *)alloc_kpage(1);
        memset(arena, 0, PAGE_SIZE);

        // 计算内存块的数量
        u32 left = PAGE_SIZE - sizeof(arena_t);
        u32 count = left / desc->size;

        // 构建 arena
        arena->count = count;
        arena->page = false;
        arena->desc = desc;

        // 分配内存块
        for (size_t i = 0; i < count; i++)
        {
            block_t *block = arena_block(i, arena, desc->size);
            list_append(&desc->free_list, &block->node);
        }
    }

    list_node_t *node = list_pop(&desc->free_list);
    block_t *block = element_entry(block_t, node, node);

    return block;
}

void kfree(void *addr)
{
    // 获得 block 和 arena
    block_t *block = (block_t *)addr;
    arena_t *arena = block_arena(block);

    // 释放内存块
    if (!arena->page)
    {
        memset(addr, 0, arena->desc->size);
        list_push(&arena->desc->free_list, &block->node);
        return;
    }
    // 整页释放
    free_kpage((u32)arena, arena->count);
}

void arena_init()
{
    u16 size = 16;
    for (size_t i = 0; i < DESC_COUNT; i++)
    {
        arena_descriptor_t *desc = &descriptors[i];
        list_init(&desc->free_list);
        desc->size = size;
        size *= 2;
    }
}